"""
1 计算体素点到点之间的MSE
  输入点云模型,计算tru_rt和pre_rt之后点云之间的mse
"""
import cupy as cp
import numpy as np
from scipy.stats import norm as gauss_norm


def norm(datas, min_v=None, max_v=None):
    if min_v is None:
        min_v = np.min(datas)
    if max_v is None:
        max_v = np.max(datas)
    return np.round((datas - min_v) / (max_v - min_v), 5)


def init_vertex(vertex_size, interval_num):
    """
    根据体素实际的大小，在其中等间隔采点。对于每个切片来说，按高斯分布，边缘的权重大，中间的权重小。
    :param vertex_size: 体素实际的大小
    :param interval_num:从体素中每条边上取多少个点，默认值的意思是会从体素中均分出100*100*100个小长方体
    :return:
    """
    # 生成小正方体中心的坐标网格
    # 使用np.mgrid生成三个一维数组，代表x, y, z坐标，然后通过reshape和transpose转换成(3, 1000)的形状
    interval = vertex_size / interval_num
    x, y, z = np.mgrid[0.5 * interval[0]:vertex_size[0]:interval[0],
              0.5 * interval[1]:vertex_size[1]:interval[1],
              0.5 * interval[2]:vertex_size[2]:interval[2]]

    # 将三维坐标网格展平为一维数组，并组合成(3, 1000)的矩阵
    centers = np.vstack((x.ravel(), y.ravel(), z.ravel())).T
    # 首先计算切片内切圆的半径
    r = np.sqrt((vertex_size[0] / 2) ** 2 + (vertex_size[1] / 2) ** 2)
    # 沿z轴方向计算每个点在xy平面到旋转中心的距离
    xy = centers[:, :2]
    center = np.array([vertex_size[0] / 2, vertex_size[1] / 2])
    distances = np.linalg.norm(xy - center[None, :], axis=1)
    weight = distances / r
    # 对权重高斯映射
    weight = gauss_norm.pdf(weight * 3)
    # 权重归一化
    weight = norm(weight)
    pt_matrix = np.zeros((centers.shape[0], 4))
    pt_matrix[:, :3] = centers
    pt_matrix[:, 3] = weight

    return pt_matrix


def cal_voxel_mse_loss(vertex_size, interval_num, rt1, rt2, mode='average', weight=1):
    # 根据体素间距，生成vm的间隔采样稀疏点云
    cube = init_vertex(vertex_size, interval_num)
    new_pos1 = np.dot(rt1, cube.T).T
    new_pos2 = np.dot(rt2, cube.T).T
    # 计算两个新位置之间顶点的距离
    distance = np.linalg.norm((new_pos1 - new_pos2), axis=1)
    distance = distance * weight
    # 给旋转中心增加权重
    if mode == 'average':
        return np.mean(distance)
    else:
        return max(distance)
